{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "af6397b2",
   "metadata": {},
   "source": [
    "# Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "839c4a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# My OpenAI Key\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"INSERT OPENAI KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d726e871",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hide INFO logs regarding token usage, etc\n",
    "import logging\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.CRITICAL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40cf0773",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "### Generate some example Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa34cd83",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import download_loader\n",
    "from llama_index.indices.vector_store import VectorStoreIndex\n",
    "from llama_index.indices.tree.base import TreeIndex\n",
    "\n",
    "WikipediaReader = download_loader(\"WikipediaReader\")\n",
    "\n",
    "loader = WikipediaReader()\n",
    "documents = loader.load_data(pages=[\"Berlin\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c32392b",
   "metadata": {},
   "source": [
    "### Create a list of any sort of indices (custom LLMs, custom embeddings, etc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f59e6c18",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [build_index_from_documents] Total LLM token usage: 0 tokens\n",
      "INFO:root:> [build_index_from_documents] Total embedding token usage: 18344 tokens\n",
      "INFO:root:> Building index from nodes: 5 chunks\n"
     ]
    }
   ],
   "source": [
    "indices = [\n",
    "    VectorStoreIndex.from_documents(documents),\n",
    "    TreeIndex.from_documents(documents),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "827ada33",
   "metadata": {},
   "source": [
    "## Using the Playground\n",
    "\n",
    "\n",
    "### Initialize with indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a04e4535",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.playground import Playground\n",
    "\n",
    "playground = Playground(indices=indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5f6999fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:openai:error_code=None error_message='Rate limit reached for default-global-with-image-limits in organization org-ehTdCqs0FpsxuTTwsJIlNSdZ on requests per min. Limit: 60.000000 / min. Current: 110.000000 / min. Contact support@openai.com if you continue to have issues. Please add a payment method to your account to increase your rate limit. Visit https://platform.openai.com/account/billing to add a payment method.' error_param=None error_type=requests message='OpenAI API error received' stream_error=False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mQuery:\u001b[0m\n",
      "What is the population of Berlin?\n",
      "\n",
      "Trying 10 combinations...\n",
      "\n",
      "\n",
      "\u001b[1mGPTVectorStoreIndex\u001b[0m, mode = default\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:openai:error_code=None error_message='Rate limit reached for default-global-with-image-limits in organization org-ehTdCqs0FpsxuTTwsJIlNSdZ on requests per min. Limit: 60.000000 / min. Current: 90.000000 / min. Contact support@openai.com if you continue to have issues. Please add a payment method to your account to increase your rate limit. Visit https://platform.openai.com/account/billing to add a payment method.' error_param=None error_type=requests message='OpenAI API error received' stream_error=False\n",
      "INFO:openai:error_code=None error_message='Rate limit reached for default-global-with-image-limits in organization org-ehTdCqs0FpsxuTTwsJIlNSdZ on requests per min. Limit: 60.000000 / min. Current: 90.000000 / min. Contact support@openai.com if you continue to have issues. Please add a payment method to your account to increase your rate limit. Visit https://platform.openai.com/account/billing to add a payment method.' error_param=None error_type=requests message='OpenAI API error received' stream_error=False\n",
      "INFO:openai:error_code=None error_message='Rate limit reached for default-global-with-image-limits in organization org-ehTdCqs0FpsxuTTwsJIlNSdZ on requests per min. Limit: 60.000000 / min. Current: 80.000000 / min. Contact support@openai.com if you continue to have issues. Please add a payment method to your account to increase your rate limit. Visit https://platform.openai.com/account/billing to add a payment method.' error_param=None error_type=requests message='OpenAI API error received' stream_error=False\n",
      "INFO:root:> [query] Total LLM token usage: 3545 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 7 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3m\n",
      "The population of Berlin in 1949 was approximately 2.2 million inhabitants. After the fall of the Berlin Wall in 1989, the population of Berlin increased to approximately 3.7 million inhabitants.\u001b[0m\n",
      "\n",
      "\u001b[1mGPTVectorStoreIndex\u001b[0m, mode = embedding\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [query] Total LLM token usage: 3545 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 7 tokens\n",
      "INFO:root:> Starting query: What is the population of Berlin?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3m\n",
      "The population of Berlin in 1949 was approximately 2.2 million inhabitants. After the fall of the Berlin Wall in 1989, the population of Berlin increased to approximately 3.7 million inhabitants.\u001b[0m\n",
      "\n",
      "\u001b[1mGPTTreeIndex\u001b[0m, mode = default\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:>[Level 0] Selected node: [1]/[1]\n",
      "INFO:root:>[Level 1] Selected node: [3]/[3]\n",
      "INFO:root:> [query] Total LLM token usage: 5168 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 0 tokens\n",
      "INFO:root:> Starting query: What is the population of Berlin?\n",
      "INFO:root:> Building index from nodes: 6 chunks\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3mThe population of Berlin is approximately 3.7 million people.\u001b[0m\n",
      "\n",
      "\u001b[1mGPTTreeIndex\u001b[0m, mode = summarize\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [query] Total LLM token usage: 21617 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 0 tokens\n",
      "INFO:root:> Starting query: What is the population of Berlin?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3m\n",
      "The population of Berlin is approximately 3.7 million people.\u001b[0m\n",
      "\n",
      "\u001b[1mGPTTreeIndex\u001b[0m, mode = embedding\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [query] Total LLM token usage: 368 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 4598 tokens\n",
      "INFO:root:> Starting query: What is the population of Berlin?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3mApproximately 3.7 million people.\u001b[0m\n",
      "\n",
      "\u001b[1mGPTTreeIndex\u001b[0m, mode = retrieve\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [query] Total LLM token usage: 1439 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 0 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3m\n",
      "The population of Berlin is 3.75 million registered inhabitants.\u001b[0m\n",
      "\n",
      "\n",
      "Ran 6 combinations in total.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Index</th>\n",
       "      <th>Mode</th>\n",
       "      <th>Output</th>\n",
       "      <th>Duration</th>\n",
       "      <th>LLM Tokens</th>\n",
       "      <th>Embedding Tokens</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>VectorStoreIndex</td>\n",
       "      <td>default</td>\n",
       "      <td>\\nThe population of Berlin in 1949 was approxi...</td>\n",
       "      <td>52.319133</td>\n",
       "      <td>3545</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>VectorStoreIndex</td>\n",
       "      <td>embedding</td>\n",
       "      <td>\\nThe population of Berlin in 1949 was approxi...</td>\n",
       "      <td>8.192025</td>\n",
       "      <td>3545</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>default</td>\n",
       "      <td>The population of Berlin is approximately 3.7 ...</td>\n",
       "      <td>12.542335</td>\n",
       "      <td>5168</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>summarize</td>\n",
       "      <td>\\nThe population of Berlin is approximately 3....</td>\n",
       "      <td>18.665586</td>\n",
       "      <td>21617</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>embedding</td>\n",
       "      <td>Approximately 3.7 million people.</td>\n",
       "      <td>3.573458</td>\n",
       "      <td>368</td>\n",
       "      <td>4598</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>TreeIndex</td>\n",
       "      <td>retrieve</td>\n",
       "      <td>\\nThe population of Berlin is 3.75 million reg...</td>\n",
       "      <td>2.269598</td>\n",
       "      <td>1439</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  Index       Mode  \\\n",
       "0  VectorStoreIndex    default   \n",
       "1  VectorStoreIndex  embedding   \n",
       "2          TreeIndex    default   \n",
       "3          TreeIndex  summarize   \n",
       "4          TreeIndex  embedding   \n",
       "5          TreeIndex   retrieve   \n",
       "\n",
       "                                              Output   Duration  LLM Tokens  \\\n",
       "0  \\nThe population of Berlin in 1949 was approxi...  52.319133        3545   \n",
       "1  \\nThe population of Berlin in 1949 was approxi...   8.192025        3545   \n",
       "2  The population of Berlin is approximately 3.7 ...  12.542335        5168   \n",
       "3  \\nThe population of Berlin is approximately 3....  18.665586       21617   \n",
       "4                  Approximately 3.7 million people.   3.573458         368   \n",
       "5  \\nThe population of Berlin is 3.75 million reg...   2.269598        1439   \n",
       "\n",
       "   Embedding Tokens  \n",
       "0                 7  \n",
       "1                 7  \n",
       "2                 0  \n",
       "3                 0  \n",
       "4              4598  \n",
       "5                 0  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "playground.compare(\"What is the population of Berlin?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8829a829",
   "metadata": {},
   "source": [
    "### Initialize with Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfbc8ade",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uses documents in a preset list of indices\n",
    "playground = Playground.from_docs(documents=documents)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
